import torch
from torch.utils.data.distributed import DistributedSampler


class MultiEpochsDataLoader(torch.utils.data.DataLoader):
    """
    扩展 PyTorch 的 DataLoader 类，使得可以在训练过程中反复迭代数据集。
    通过将数据加载器与 `_RepeatSampler` 结合，能够支持无穷次循环迭代数据集。
    """

    def __init__(self, *args, **kwargs):
        """
        初始化 MultiEpochsDataLoader，使用 _RepeatSampler 来实现数据的重复采样。

        参数：
            *args: 传递给 DataLoader 的参数。
            **kwargs: 传递给 DataLoader 的关键字参数。
        """
        super().__init__(*args, **kwargs)
        # 标记 DataLoader 是否已初始化
        self._DataLoader__initialized = False
        # 将原始的 batch_sampler 替换为 _RepeatSampler 实例，实现数据的无限循环
        self.batch_sampler = _RepeatSampler(self.batch_sampler)
        # 标记 DataLoader 初始化完成
        self._DataLoader__initialized = True
        # 使用原 DataLoader 的 __iter__ 方法，初始化数据迭代器
        self.iterator = super().__iter__()

    def __len__(self):
        """
        获取数据集的长度（即样本数量）。
        返回：batch_sampler 的样本数量
        """
        return len(self.batch_sampler.sampler)

    def __iter__(self):
        """
        返回数据的迭代器，确保每次迭代返回下一个 batch。
        使用 _RepeatSampler 使得迭代器能够在每个 epoch 结束后继续采样。
        """
        for i in range(len(self)):
            # 每次返回一个新的 batch
            yield next(self.iterator)


class _RepeatSampler(object):
    """
    一个重复采样器（Sampler），用于支持在训练中无休止地重复数据集。

    参数：
        sampler (Sampler): 传入的原始 sampler，通常是 DataLoader 使用的 batch_sampler。
    """

    def __init__(self, sampler):
        # 初始化时接受一个原始的 sampler 对象
        self.sampler = sampler

    def __iter__(self):
        """
        无限循环迭代原始采样器的样本数据，直到手动停止。
        """
        while True:
            # 使用 yield from 来重复地从原始采样器中获取样本
            yield from iter(self.sampler)

